﻿#include "asycipc/frame.hpp"

#include <algorithm>
#include <array>
#include <cstring>
#include <string>

#include "asycipc/crc32c.hpp"

namespace asycipc {
namespace {
void WriteUint16(std::vector<uint8_t>& out, uint16_t value) {
  out.push_back(static_cast<uint8_t>((value >> 8) & 0xFFu));
  out.push_back(static_cast<uint8_t>(value & 0xFFu));
}

void WriteUint32(std::vector<uint8_t>& out, uint32_t value) {
  out.push_back(static_cast<uint8_t>((value >> 24) & 0xFFu));
  out.push_back(static_cast<uint8_t>((value >> 16) & 0xFFu));
  out.push_back(static_cast<uint8_t>((value >> 8) & 0xFFu));
  out.push_back(static_cast<uint8_t>(value & 0xFFu));
}

void WriteUint64(std::vector<uint8_t>& out, uint64_t value) {
  for (int shift = 56; shift >= 0; shift -= 8) {
    out.push_back(static_cast<uint8_t>((value >> shift) & 0xFFu));
  }
}

uint16_t ReadUint16(const std::vector<uint8_t>& data, std::size_t offset) {
  return static_cast<uint16_t>(data[offset] << 8 | data[offset + 1]);
}

uint32_t ReadUint32(const std::vector<uint8_t>& data, std::size_t offset) {
  return (static_cast<uint32_t>(data[offset]) << 24) |
         (static_cast<uint32_t>(data[offset + 1]) << 16) |
         (static_cast<uint32_t>(data[offset + 2]) << 8) |
         static_cast<uint32_t>(data[offset + 3]);
}

uint64_t ReadUint64(const std::vector<uint8_t>& data, std::size_t offset) {
  uint64_t value = 0;
  for (int i = 0; i < 8; ++i) {
    value = (value << 8) | static_cast<uint64_t>(data[offset + i]);
  }
  return value;
}

constexpr std::size_t kHeaderPrefixSize = 18;  // magic..payloadLength
constexpr std::size_t kHeaderSize = 24;        // full header including checksum/ext
constexpr std::size_t kTrailerSize = 4;        // payload checksum
}  // namespace

std::vector<uint8_t> Frame::Serialize() const {
  Frame frame = *this;
  frame.magic = kFrameMagic;
  frame.version = kProtocolVersion;
  frame.payload_length = static_cast<uint32_t>(frame.payload.size());
  if (frame.extensions.size() > 0xFFFFu) {
    throw FrameValidationError("extension too large");
  }
  frame.ext_length = static_cast<uint16_t>(frame.extensions.size());

  std::vector<uint8_t> prefix;
  prefix.reserve(kHeaderPrefixSize);
  WriteUint16(prefix, frame.magic);
  prefix.push_back(frame.version);
  prefix.push_back(static_cast<uint8_t>(frame.type));
  prefix.push_back(frame.flags);
  prefix.push_back(frame.reserved);
  WriteUint64(prefix, frame.seq_no);
  WriteUint32(prefix, frame.payload_length);

  frame.header_checksum = Crc32c(prefix.data(), prefix.size());

  std::vector<uint8_t> buffer;
  buffer.reserve(kHeaderSize + frame.extensions.size() + frame.payload.size() +
                 kTrailerSize);
  buffer.insert(buffer.end(), prefix.begin(), prefix.end());
  WriteUint32(buffer, frame.header_checksum);
  WriteUint16(buffer, frame.ext_length);
  buffer.insert(buffer.end(), frame.extensions.begin(), frame.extensions.end());

  if (!frame.payload.empty()) {
    frame.payload_checksum = Crc32c(frame.payload.data(), frame.payload.size());
  } else {
    frame.payload_checksum = Crc32c(static_cast<const uint8_t*>(nullptr), 0);
  }

  buffer.insert(buffer.end(), frame.payload.begin(), frame.payload.end());
  WriteUint32(buffer, frame.payload_checksum);
  return buffer;
}

Frame Frame::Deserialize(const std::vector<uint8_t>& raw) {
  if (raw.size() < kHeaderSize + kTrailerSize) {
    throw FrameValidationError("frame too small");
  }
  Frame frame;
  frame.magic = ReadUint16(raw, 0);
  frame.version = raw[2];
  frame.type = static_cast<FrameType>(raw[3]);
  frame.flags = raw[4];
  frame.reserved = raw[5];
  frame.seq_no = ReadUint64(raw, 6);
  frame.payload_length = ReadUint32(raw, 14);
  frame.header_checksum = ReadUint32(raw, 18);
  frame.ext_length = ReadUint16(raw, 22);

  if (frame.magic != kFrameMagic) {
    throw FrameValidationError("invalid magic");
  }
  if (frame.version != kProtocolVersion) {
    throw FrameValidationError("unsupported version");
  }

  std::vector<uint8_t> prefix(raw.begin(), raw.begin() + kHeaderPrefixSize);
  const uint32_t expected = Crc32c(prefix.data(), prefix.size());
  if (expected != frame.header_checksum) {
    throw FrameValidationError("header checksum mismatch");
  }

  const std::size_t extensions_offset = kHeaderSize;
  if (raw.size() < extensions_offset + frame.ext_length + frame.payload_length +
                      kTrailerSize) {
    throw FrameValidationError("incomplete frame");
  }

  frame.extensions.assign(raw.begin() + extensions_offset,
                          raw.begin() + extensions_offset + frame.ext_length);

  const std::size_t payload_offset = extensions_offset + frame.ext_length;
  frame.payload.assign(raw.begin() + payload_offset,
                       raw.begin() + payload_offset + frame.payload_length);
  frame.payload_checksum = ReadUint32(raw, payload_offset + frame.payload_length);

  const uint32_t computed_payload_crc =
      frame.payload.empty()
          ? Crc32c(static_cast<const uint8_t*>(nullptr), 0)
          : Crc32c(frame.payload.data(), frame.payload.size());
  if (computed_payload_crc != frame.payload_checksum) {
    throw FrameValidationError("payload checksum mismatch");
  }

  return frame;
}

Frame Frame::MakeData(uint64_t seq, const std::vector<uint8_t>& payload,
                      const std::vector<uint8_t>& extensions, uint8_t flags) {
  Frame frame;
  frame.type = FrameType::kData;
  frame.seq_no = seq;
  frame.flags = flags;
  frame.extensions = extensions;
  frame.payload = payload;
  return frame;
}

Frame Frame::MakeAck(uint64_t seq) {
  Frame frame;
  frame.type = FrameType::kAck;
  frame.seq_no = seq;
  return frame;
}

Frame Frame::MakeNack(uint64_t seq, const std::string& reason) {
  Frame frame;
  frame.type = FrameType::kNack;
  frame.seq_no = seq;
  if (!reason.empty()) {
    frame.payload.assign(reason.begin(), reason.end());
  }
  return frame;
}

Frame Frame::MakeHeartbeat() {
  Frame frame;
  frame.type = FrameType::kHeartbeat;
  frame.seq_no = 0;
  return frame;
}

Frame Frame::MakeHeartbeatAck() {
  Frame frame;
  frame.type = FrameType::kHeartbeatAck;
  frame.seq_no = 0;
  return frame;
}

}  // namespace asycipc
